#include <ntddk.h>
#include <wdm.h>

typedef unsigned int UINT;
typedef char * PCHAR;

#pragma pack(push, 1)    
struct SGDT_Result
{
        UINT16 size;
        UINT32 base;
};
#pragma pack(pop)


NTSTATUS Close(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS WriteBufferedIO(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS ReadBufferedIO(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS UnSupportedFunction(PDEVICE_OBJECT DeviceObject, PIRP Irp);
VOID Unload(PDRIVER_OBJECT  DriverObject); 

NTSTATUS DriverEntry(PDRIVER_OBJECT  pDriverObject, PUNICODE_STRING  pRegistryPath)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    UINT uiIndex = 0;
    PDEVICE_OBJECT pDeviceObject = NULL;
    UNICODE_STRING usDriverName, usDosDeviceName;
    ULONG DeviceExtensionSize;

    DbgPrint("DriverEntry Called \r\n");

    RtlInitUnicodeString(&usDriverName, L"\\Device\\MyDevice");
    RtlInitUnicodeString(&usDosDeviceName, L"\\DosDevices\\MyDevice"); 

    NtStatus = IoCreateDevice(pDriverObject, 0, &usDriverName, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, FALSE, &pDeviceObject);

    if(NtStatus == STATUS_SUCCESS)
    {

        for(uiIndex = 0; uiIndex < IRP_MJ_MAXIMUM_FUNCTION; uiIndex++)
             pDriverObject->MajorFunction[uiIndex] = UnSupportedFunction;
    
        pDriverObject->MajorFunction[IRP_MJ_CLOSE]             = Close;
        pDriverObject->MajorFunction[IRP_MJ_READ]              = ReadBufferedIO;
    
        pDriverObject->DriverUnload = Unload;

        pDeviceObject->Flags |= DO_BUFFERED_IO;
    
        pDeviceObject->Flags &= (~DO_DEVICE_INITIALIZING);
    
        IoCreateSymbolicLink(&usDosDeviceName, &usDriverName);
    }


    return NtStatus;
}

VOID Unload(PDRIVER_OBJECT  DriverObject)
{    
    
    UNICODE_STRING usDosDeviceName;
    
    DbgPrint("Unload Called \r\n");
    
    RtlInitUnicodeString(&usDosDeviceName, L"\\DosDevices\\MyDevice");
    IoDeleteSymbolicLink(&usDosDeviceName);

    IoDeleteDevice(DriverObject->DeviceObject);
}

NTSTATUS Close(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    DbgPrint("Close Called \r\n");

    return NtStatus;
}



NTSTATUS ReadBufferedIO(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_BUFFER_TOO_SMALL;
    PIO_STACK_LOCATION pIoStackIrp = NULL;
    PCHAR pReadDataBuffer;
    ULONG messageSize = 0;
    struct SGDT_Result sgdtResult;


    DbgPrint("ReadBufferedIO Called \r\n");
    
    pIoStackIrp = IoGetCurrentIrpStackLocation(Irp);
    
    if(pIoStackIrp)
    {
        messageSize = pIoStackIrp->Parameters.Read.Length;
        pReadDataBuffer = (PCHAR)Irp->AssociatedIrp.SystemBuffer;
    
        if(pReadDataBuffer)
        {     
            __asm { 
                SGDT sgdtResult;
            }   
            DbgPrint("GDT size = %u \n", sgdtResult.size);
            DbgPrint("GDT base = %u \n", sgdtResult.base);                    
            RtlCopyMemory(pReadDataBuffer, sgdtResult.base , sgdtResult.size + 1);
            NtStatus = STATUS_SUCCESS;
        }
    }


    Irp->IoStatus.Status = NtStatus;
    Irp->IoStatus.Information = messageSize;

    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return NtStatus;
}

NTSTATUS UnSupportedFunction(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_NOT_SUPPORTED;
    DbgPrint("UnSupportedFunction Called \r\n");

    return NtStatus;
}